SHMEM UDMA Programming

导言

原生 URMA 要显式管理 Context、Segment、Jetty、WR 和完成队列。cann/shmem 提供了更高层的 PGAS 编程模型：程序主要面对“对称地址、目标 PE、put/get 和同步”。本文先解释 aclshmem_my_pe()、aclshmem_n_pes() 和 aclshmemx_udma_put_nbi()，再用两 PE 的 put-with-signal Case 说明数据和通知如何配合。

系列位置

1. URMA Mental Model：对象与数据路径
2. URMA Write and Read：最小读写程序
3. URMA Completion and Concurrency：完成、顺序与并发
4. SHMEM UDMA Programming：对称内存与 put 接口
5. AIV UDMA Direct Drive：st_dev、多 QP 与 Relay

源码截面

本文的 UDMA 扩展签名与并发限制核对到 cann/shmem@59f0313325ef3179306ce6fa3bfef8ba2ff4cbc5。这类新接口仍在演进，使用其他 revision 时应重新检查公共头文件和示例。

从 URMA 到 SHMEM

两层接口解决的是同一类远端内存访问，但暴露的抽象不同：

原生 URMA	cann/shmem	封装后的变化
本地、远端 Segment 句柄	对称堆中的指针	库维护远端内存映射
Jetty 或目标端点	PE 编号	用逻辑参与者选择目标
显式构造 WR/SGE	put/get 参数	库生成底层传输请求
轮询 JFC	quiet、signal、wait 等	用 SHMEM 同步语义表达完成和发布
应用管理队列资源	初始化阶段建立通信资源	数据面调用更短

所以，aclshmem_put... 并不要求用户先手写 urma_jfs_wr_t。它把底层资源和 WQE 构造藏在了库内部。

PE 是什么

PE，Processing Element，是 SHMEM 世界中的逻辑通信参与者。 在常见部署中，一个进程或一张卡对应一个 PE，编号范围为 0 到 n_pes - 1。

int my_pe = aclshmem_my_pe();
int n_pes = aclshmem_n_pes();

• **aclshmem_my_pe()**：返回当前代码属于哪个 PE，可类比通信库中的 rank。
• **aclshmem_n_pes()**：返回本次 SHMEM 作业一共有多少个 PE，可类比 world size。

例如，两进程启动时：

进程	aclshmem_my_pe()	aclshmem_n_pes()
进程 A	0	2
进程 B	1	2

PE 不是 AIV 核编号。一个 PE 内仍然可以启动多个 AIV block；AIV 使用 GetBlockIdx() 区分，PE 使用 aclshmem_my_pe() 区分。

对称内存是什么

所有 PE 要以相同顺序和大小从对称堆分配对象。由此，每个 PE 都能用相同的“对称地址表达”指定另一个 PE 上对应的对象。

PE 0 对称堆：payload(offset X)  signal(offset Y)
PE 1 对称堆：payload(offset X)  signal(offset Y)

当 PE 0 调用：

aclshmem_putmem(dst, src, bytes, 1);

dst 在代码中表现为本地可计算的对称指针，最后一个参数 1 才指定实际目标是 PE 1 上相同对称偏移的对象。

对称不等于内容相同

对称约束的是分配布局和寻址规则，不要求各 PE 的对象内容相同，也不要求所有 PE 执行完全相同的数据操作。

aclshmemx_udma_put_nbi 逐参数解释

显式 UDMA 扩展接口可写成：

aclshmemx_udma_put_nbi<T>(
    dst,
    src,
    buf,
    elem_size,
    pe,
    sync_id);

函数名可以拆成六部分：

• **aclshmemx**：扩展接口，而不是基础 SHMEM API。
• **udma**：明确选择 UDMA 传输引擎。
• **put**：从本地向远端写。
• **nbi**：non-blocking immediate，调用返回不代表传输已完成。
• **<T>**：模板元素类型。
• 无 qp 后缀：使用默认选择的 QP；多 QP 版本会显式接收 qp_idx。

参数含义如下：

参数	所在内存/含义	初学者最容易错的地方
dst	目标对象的对称地址表达	它不是普通远端裸指针；真正目标由 pe 决定
src	当前 PE 的本地 GM 源地址	nbi 完成前不要覆盖或释放
buf	AIV 提供的 UB 临时区	默认 PIPE_MTE3 路径用它暂存 WQE，至少提供 ACLSHMEM_UDMA_MTE_STAGING_UB_SIZE 字节；它不是 payload 目标
elem_size	T 类型的元素数量	不要在模板类型非字节时误传字节数
pe	目标 PE 编号	合法范围是 0 到 n_pes - 1
sync_id	Device 流水同步事件编号	默认 PIPE_MTE3 路径使用；选择 PIPE_S 时被忽略

当前源码截面规定单次显式 UDMA put 最多传输 256 MB，实际字节数为 elem_size * sizeof(T)。更大消息需要切成多条请求；升级版本后仍应重新检查公共头文件。

调用后需要建立完成边界：

aclshmemx_udma_put_nbi<uint8_t>(
    dst, src, wqe_ub, bytes, peer, sync_id);

// 在复用 src、wqe_ub 或依赖远端结果前等待相应 UDMA 操作完成。
aclshmemx_udma_quiet(peer);

具体 quiet 接口是否按 PE、按 QP 或按全局生效，应以选择的 UDMA API 变体为准。

默认 UDMA 接口不能被多核随意共享

当前公共头明确写明：使用不带 qp_idx 的 UDMA RMA/AMO 接口时，不支持对同一个 PE 并发操作。需要多 AIV 并行时，应先配置多个 QP，再让每个 AIV 使用不同 qp_idx 的接口。

SHMEM 暴露了哪些接口

API 很多时，按“操作、完成、引擎”三个维度分类比背函数名有效。

通用 RMA

接口族	用途
aclshmem_putmem()	按字节把本地数据写到远端
aclshmem_getmem()	按字节把远端数据读到本地
aclshmem_<type>_put/get()	以具体数据类型和元素数操作
aclshmem_*_put_nbi/get_nbi()	非阻塞发起 put/get

通用接口可以由库根据拓扑和配置选择 MTE、SDMA、UDMA 或 RDMA 等数据路径。调用者表达的是“我要做 RMA”，不一定固定某个引擎。

数据加通知

接口族	用途
aclshmem_putmem_signal()	写数据后更新远端 signal
aclshmem_putmem_signal_nbi()	非阻塞发起写数据加 signal
aclshmem_<type>_put_signal()	typed 版本的 put-with-signal
aclshmem_signal_wait_until()	等待当前 PE 的本地 signal 满足条件
aclshmemx_signal_op()	对远端 signal 执行 SET/ADD

putmem_signal 的价值不只是少调用一个函数，而是把“数据先到，signal 后发布”表达成一个有顺序保证的操作。

完成与顺序

接口	核心问题
aclshmem_quiet()	此前由当前 PE 发起的 SHMEM 操作是否完成
aclshmem_fence()	此前操作与后续操作之间如何排序
aclshmem_barrier() / aclshmem_barrier_all()	多个 PE 是否都到达同步点，并满足相应完成语义

在 cann/shmem 的某些固定实现版本中，aclshmem_fence() 可能使用与 quiet 相同的强实现。但编写可移植逻辑时仍应区分：Fence 主要表达顺序，Quiet 主要表达完成。

显式 UDMA 扩展

公开 UDMA 头文件还提供以下接口族：

• aclshmemx_udma_put_nbi()、aclshmemx_udma_get_nbi()。
• aclshmemx_udma_qp_put_nbi()、aclshmemx_udma_qp_get_nbi()：显式指定 qp_idx。
• aclshmemx_udma_relay_put_nbi()、对应 Get 版本：通过 relay_pe 选择中继路径。
• aclshmemx_udma_qp_put_signal_nbi() 等 UDMA put-with-signal 版本。
• aclshmemx_udma_quiet(pe)、aclshmemx_udma_qp_quiet(pe, qp_idx)：按 PE 或指定 QP 等待完成。
• UDMA 原子操作和 QP 数量配置接口。

这些扩展让调用者直接控制 UDMA 数据面，但同时也承担更多 QP 所有权、UB 临时区、完成和同步责任。

两 PE 的数据发布 Case

目标：PE 0 把 payload 写给 PE 1，PE 1 只有在数据可用后才读取。

对称资源

Host 侧需要让所有 PE 以相同顺序分配：

// 教学骨架：初始化、错误检查和实际类型按对应版本补齐。
uint8_t *payload = (uint8_t *)aclshmem_malloc(payload_bytes);
int32_t *data_ready = (int32_t *)aclshmem_malloc(ACLSHMEM_SIGNAL_SIZE);

两个 PE 都执行这两次分配，即使只有 PE 1 最终消费远端 payload。

Device 侧协议

int pe = aclshmem_my_pe();
int npes = aclshmem_n_pes();

if (npes != 2 || GetBlockIdx() != 0) {
    return;
}

if (pe == 0) {
    fill_local_payload(payload);

    aclshmem_putmem_signal(
        payload,                 // PE 1 上对应的对称 payload
        payload,                 // PE 0 的本地源
        payload_bytes,
        data_ready,              // PE 1 上对应的 signal
        1,                       // signal 值
        ACLSHMEM_SIGNAL_SET,
        1);                      // 目标 PE
} else {
    aclshmem_signal_wait_until(
        data_ready,
        ACLSHMEM_CMP_EQ,
        1);

    consume_local_payload(payload);
}

逐步发生的事情是：

1. PE 0 的 AIV 填好本地 payload。
2. putmem_signal 把 payload 写到 PE 1 的对称地址。
3. 同一个有序操作再把 PE 1 的 data_ready 设置为 1。
4. PE 1 的 wait 轮询的是自己的本地 signal 地址。
5. wait 满足后，PE 1 才读取自己的 payload。

这解释了“发送端传的是 data_ready，接收端为什么也拿本地 data_ready 等待”：对称指针表达的是相同偏移，pe 决定实际远端实例。

wait 先执行还是后执行

两种时序都应工作：

情况 A：PE 1 先 wait ───────────────→ signal 到达后退出
情况 B：signal 已经是 1 → PE 1 再 wait → 立即满足

它等待的是内存中的条件，而不是只能接住一次的瞬时通知。

多轮不能一直等待 == 1

若第二轮仍把 signal 设置为 1，PE 1 可能直接看到上一轮留下的值。更稳妥的做法是使用递增序号：

// 第 round 轮，round 从 1 开始。
uint64_t expected = round;

// 发送端发布 expected。
// 接收端等待 signal >= expected。
aclshmem_signal_wait_until(signal, ACLSHMEM_CMP_GE, expected);

另一种做法是复位后握手，但必须保证发送端不会在接收端复位前开始下一轮。多生产者场景还应使用独立 signal，避免不同写者互相覆盖。

nbi 版本如何配合

aclshmem_putmem_signal_nbi() 适合让发起与后续计算重叠，但返回后不能立即复用相关资源：

aclshmem_putmem_signal_nbi(...);

do_independent_compute();

aclshmem_quiet();
reuse_source_buffer();

接收端仍可通过 aclshmem_signal_wait_until() 建立消费边界。发送端的 quiet 解决本端完成，接收端的 wait 解决远端何时可消费，两者站在不同参与者的视角。

cann/shmem 内部如何接到 UDMA

不能简单理解为：

aclshmem_put → Host 调用 urma_post_jetty_send_wr

更接近实际的分层是：

Host 控制面
  HCOMM 创建 Endpoint、注册内存、创建 Channel
  获取 SQ/CQ、Doorbell、远端内存等描述
  把 PE/QP 表和设备状态下发给 Device

Device 数据面
  AIV 根据 pe/qp 找到资源
  把对称地址换算成目标地址
  检查队列信用
  在 UB 中组装 UDMA WQE
  把 WQE 写入 HBM SQ
  更新 head 并敲 Doorbell
  统计或轮询 CQ 完成

也就是说，SHMEM 在 Host 初始化阶段准备资源，在 AIV 执行期间直接驱动数据面。应用少写了 URMA 对象管理代码，但底层的 Segment、队列、WQE、Doorbell 和完成机制仍然存在。

本篇结论

• PE 是逻辑通信参与者，不是 QP，也不是 AIV block。
• 对称指针给出远端对象的偏移坐标，pe 选择目标实例。
• nbi 只表示非阻塞发起，必须再建立完成边界。
• put-with-signal 把数据传输和发布通知按正确顺序组合起来。
• wait 在接收端轮询本地 signal；它与发送端指向的远端对称 signal 是同一逻辑对象。
• SHMEM 封装了 URMA 风格资源管理，但没有消除底层队列和 WQE。

参考资料

• CANN SHMEM 源码仓
• UDMA Device 公共头文件
• Signal Operations 公共头文件
• Device Team 公共头文件
• 深入对称寻址：SHMEM Symmetric Memory
• 完整单向教学工程：AIV Direct Drive Programming